from pydantic import BaseModel
from langchain.llms import Ollama
from langchain_core.prompts import PromptTemplate
from config import OLLAMA_BASE_URL, LLM_MODEL

# 定义Pydantic模型，用于结构化输出
class AnswerModel(BaseModel):
    summary: str
    details: str

# 定义提示模板，引导模型生成符合结构的回答
template = """
请根据以下问题，生成JSON格式的回答，包含字段："summary"（摘要），"details"（详细信息）。
问题: {question}
回答:
"""
prompt = PromptTemplate.from_template(template)

# 初始化LLM
llm = Ollama(base_url=OLLAMA_BASE_URL, model=LLM_MODEL)

# 生动形象的测试问题
question = "请用一个生动的比喻解释什么是LangChain？"

# 构造完整提示
input_prompt = prompt.format(question=question)

# 调用模型生成回答
response_text = llm.invoke(input_prompt)

print("模型原始输出:")
print(response_text)

# 使用Pydantic解析模型输出
try:
    answer = AnswerModel.parse_raw(response_text)
    print("\n结构化解析结果:")
    print(f"摘要: {answer.summary}")
    print(f"详细信息: {answer.details}")
except Exception as e:
    print(f"解析失败: {e}")